/******************************************************************************
** @file	WMIReaderTest.cpp
** @author  Arturo Aguilar
** @author  Andrea Herbas
**
** @date	Aug 1, 2010
**
** @desc	This class test all the dll functionallity. 
**			
**          - tests processes running in the operating system 
**          - tests the default shares 
**          - tests some default diretories and files 
**          - tests the registries 
**			
**
******************************************************************************/

#include <iostream>
#include <vector> 
#include <string>
#include <algorithm>
#include <fstream>
#include <vector>
//#include "../../wmi_filedirmanager.h"
#include "../Unit--/unit--.h"
#include <wmi_reader.h>

using namespace std;
using namespace unit_minus;



int* active;
bool localActive = true;
bool shareActive = false;
bool registryActive = false;
bool fileActive = false;
bool processActive = false;
bool driveActive = false;

wstring user;
wstring password;
wstring host;

vector<string> sharedList;
vector<string> fileList;
vector<string> registryList;
vector<string> processList;
vector<string> driveList;


void
getCommands(const wstring& input,wstring& command,wstring& result)
{	
    int name = input.rfind(L"=", input.size() - 1);
    result = input.substr(name + 1);
	command = input.substr(0,name);	
}

void 
openConfig()
{
	if(active != NULL)
	{		
		return;
	}

	active=new int(1);

	bool band=true;

	ifstream f("TestConfig.txt");

    string buffer;
    while (!f.eof())
    {
        getline(f, buffer);
		wstring temp(buffer.length(),L' ');
		copy(buffer.begin(), buffer.end(), temp.begin());
		
		if(temp.size() >= 2)
		{
			if(temp[0] !=  BACKSLASH && temp[1] != BACKSLASH)
			{
				wstring item, result;
				getCommands(temp,item,result);

                string res(result.begin(), result.end());

				if(item.compare(L"share") == 0)
				{
					sharedList.push_back(res);	
				}

				else if(item.compare(L"registryPath") == 0)
				{
					registryList.push_back(res);
				}

				else if(item.compare(L"file") == 0)
				{
					fileList.push_back(res);
				}

				else if(item.compare(L"process") == 0)
				{
					processList.push_back(res);
				}

				else if(item.compare(L"drive") == 0)
				{
					driveList.push_back(res);
				}

				if(item.compare(L"driveActive") == 0)
				{
					if(result.compare(L"true") == 0)
					{
						driveActive = true;
					}
				}

				else if(item.compare(L"shareActive") == 0)
				{
					if(result.compare(L"true") == 0)
					{
						shareActive = true;
					}
				}

				else if(item.compare(L"registryActive") == 0)
				{
					if(result.compare(L"true") == 0)
					{
						registryActive = true;
					}
				}

				else if(item.compare(L"fileActive") == 0)
				{
					if(result.compare(L"true") == 0)
					{
						fileActive = true;
					}
				}

				else if(item.compare(L"processActive") == 0)
				{
					if(result.compare(L"true") == 0)
					{
						processActive = true;
					}
				}

				else if(item.compare(L"user") == 0)
				{
					user=result;
				}

				else if(item.compare(L"password") == 0)
				{
					password=result;
				}
				else if(item.compare(L"host") == 0)
				{
					host=result;
				}	
			}
		}	
    }
    f.close();

	if(!localActive)
	{
		wcout<<"USER="<<user<<endl;;
		wcout<<"PASSWORD="<<password<<endl;
		wcout<<"MACHINE="<<host<<endl;
	}
	else
	{
		wcout<<"LocalActive = true"<<endl;
	}

	wcout<<"shareActive = "<<shareActive<<endl;
	wcout<<"registryActive = "<<registryActive<<endl;
	wcout<<"fileActive = "<<fileActive<<endl;
	wcout<<"processActive = "<<processActive<<endl;
	wcout<<"driveActive = "<<driveActive<<endl;

	system("PAUSE");
}


WMI_HResult* 
initConnection(WMI_Connection& connection, bool connectionType)
{ 
	openConfig();


	WMI_ConnectionManager* wmiMng = new WMI_ConnectionManager();
	WMI_HResult* whres = wmiMng->initCOMSecurityConfiguration();

	if(connectionType)
	{
		if(!localActive)
		{			
			whres = wmiMng->createConnection(user, password, host, PC_NAMESPACE_REG, connection);	
		}
		else
		{
			whres = wmiMng->createConnection(PC_LOCAL_USER, PC_LOCAL_PASSWORD, PC_NAMESPACE_REG, connection);	
		}
	}
	else
	{
		if(!localActive)
		{
			whres = wmiMng->createConnection(user, password, host, PC_NAMESPACE, connection);	
		}
		else
		{
			whres = wmiMng->createConnection(PC_LOCAL_USER, PC_LOCAL_PASSWORD, PC_NAMESPACE, connection);	
		}
	}
	return whres;
}

testSuite(WMIReaderSuite);

testCase (getSharesTest, WMIReaderSuite)
{
	cout << "---------- Start Share Test ----------" << endl;

	while(true)
	{
		WMI_Connection wmiConnection;
		WMI_HResult* whres = initConnection(wmiConnection, false);

		if(FAILED(whres->getHRESULT() ))
		{
			assertTrue(!FAILED(whres->getHRESULT()));
			break;
		}

		whres = wmiConnection.executeQuery(L"SELECT * FROM Win32_Share");
		if( FAILED(whres->getHRESULT()))
		{
			assertTrue(!FAILED(whres->getHRESULT()));
			break;
		}

		IEnumWbemClassObject* pEnumerator = NULL;
		pEnumerator = wmiConnection.getIEnumWbemClassObject();
	   
		WMI_Result<WMI_Share> listShares;
		listShares.fillVector(pEnumerator);

		if(!shareActive)
        {
            sharedList.clear();
            sharedList.push_back(string(C_SHARE.begin(), C_SHARE.end()));
            sharedList.push_back(string(IPC_SHARE.begin(), IPC_SHARE.end()));
            sharedList.push_back(string(ADMIN_SHARE.begin(), ADMIN_SHARE.end()));
        }

		int shareCount = 0;

        for(int i = 0; i < sharedList.size(); i++)
		{
			for(int j = 0; j < listShares.getObjects().size(); j++)
			{
                wstring share(sharedList[i].begin(), sharedList[i].end());
				if(share.compare((listShares.getObjects()[j])->getName()) == 0)
				{
                    assertTrue(share.compare((listShares.getObjects()[j])->getName()) == 0);
					shareCount++;
				}
			}	
		}
        listShares.showObjects();
		wmiConnection.close();

        assertTrue(shareCount == sharedList.size());

		cout << "---------- End Share Test ----------" << endl << endl;

		system("PAUSE");
		break;
		}
}

testCase (getProccessesTest, WMIReaderSuite)
{
	cout << "---------- Start Processes Test ----------" << endl;

    while(true)
	{
		WMI_Connection wmiConnection;
		WMI_HResult* whres = initConnection(wmiConnection, false);

		if(FAILED(whres->getHRESULT() ))
		{
			assertTrue(!FAILED(whres->getHRESULT()));
			break;
		}

		whres = wmiConnection.executeQuery(L"SELECT * FROM Win32_Process");

		if( FAILED(whres->getHRESULT()))
		{
			assertTrue(!FAILED(whres->getHRESULT()));
			break;
		}

		IEnumWbemClassObject* pEnumerator = NULL;
		pEnumerator = wmiConnection.getIEnumWbemClassObject();
	   
		WMI_Result<WMI_Process> listProcess;
		listProcess.fillVector(pEnumerator);

		int processCount = 0;

        if(!processActive)
        {
            processList.clear();
            processList.push_back(string(EXPLORER.begin(), EXPLORER.end()));
            processList.push_back(string(SVCHOST.begin(), SVCHOST.end()));
            processList.push_back(string(SPOOLSV.begin(), SPOOLSV.end()));
        }

		for(int j = 0; j < listProcess.getObjects().size(); j++)
		{
            for(int i = 0; i < processList.size(); i++)
            {
                wstring process(processList[i].begin(), processList[i].end());
                if(process.compare((listProcess.getObjects()[j])->getName()) == 0)
				{
					assertTrue(process.compare((listProcess.getObjects()[j])->getName()) == 0);
				}
            }
		}
		listProcess.showObjects();

		wmiConnection.close();
		break;
	}
	cout << "---------- End Processes Test ----------" << endl << endl;
	system("PAUSE");
}

testCase (WMIReaderRegistryTest, WMIReaderSuite)
{
	cout << "---------- Start Registry children Test ----------" << endl;

	WMI_Connection wmiConnection;
	
	WMI_HResult* whres = initConnection(wmiConnection, true);

	if(FAILED(whres->getHRESULT() ))
	{
		assertTrue(!FAILED(whres->getHRESULT()));
	}

	WMI_RegistryManager man(&wmiConnection);
	WMI_Registry* rclass = new WMI_Registry();

	int registryNum = 0;

    if(!registryActive)
    {
        registryList.clear();
        registryList.push_back(string(TSAM.begin(), TSAM.end()));
        registryList.push_back(string(TSYSTEM.begin(), TSYSTEM.end()));
        registryList.push_back(string(TSECURITY.begin(), TSECURITY.end()));
        registryList.push_back(string(TSOFTWARE.begin(), TSOFTWARE.end()));
    }

    wstring root;
    wstring subkey;

    for(int j = 0; j < registryList.size(); j++)
    {
        wstring path(registryList[j].begin(), registryList[j].end());

        man.getParseRoot(path, root, subkey);

        HiveEnum rootEnum = man.convertStringToHKEY(root);
        wstring subKeyReg;
        wstring outPath;
        man.getRegistryName(subkey, outPath, subKeyReg);
        
        if(!subKeyReg.empty() && !outPath.empty())
        { 
            man.getRegistrys(rootEnum, outPath, &rclass);

            for(int i = 0; i < rclass->getRegistryList().size(); i++ )
	        {   
		        WMI_Registry* aux = rclass->getRegistryList()[i];
               
                wstring outRegistryName = aux->getName();

                transform(outRegistryName.begin(), outRegistryName.end(), outRegistryName.begin(), towlower);
                transform(subKeyReg.begin(), subKeyReg.end(), subKeyReg.begin(), towlower);

                if(subKeyReg.compare(outRegistryName) == 0)
                {
                    assertTrue(subKeyReg.compare(outRegistryName) == 0);
		            registryNum++;
                }  
	        }
        }
        else
        {
            registryList.clear();
        }
    } 
	
	WMI_Result<WMI_Registry> registry;
	registry.setWRegistry(rclass);

    wmiConnection.close();

    assertTrue(registryNum >= registryList.size());

    cout << "---------- End Registry children Test ----------" << endl << endl;
    system("PAUSE");
}

testCase (WMIReaderStandartRegistriesTest, WMIReaderSuite)
{
	cout << "---------- Start Standart Registries Test ----------" << endl << endl;

	WMI_Connection wmiConnection;
	
	WMI_HResult* whres = initConnection(wmiConnection, true);

	if(FAILED(whres->getHRESULT() ))
	{
		assertTrue(!FAILED(whres->getHRESULT()));
	}

	WMI_RegistryManager man(&wmiConnection);

	vector<WMI_Registry> rclass;

	man.getStandardRegistryHives(rclass);

	int registryNum = 1;

	for(int i = 0; i < rclass.size(); i++ )
	{
     	WMI_Registry* aux = &rclass[i];

		if(aux->getName().compare(CONST_HKEY_LOCAL) == 0)
		{
			assertTrue(aux->getName().compare(CONST_HKEY_LOCAL) == 0);
			registryNum++;
		}
		else if(aux->getName().compare( CONST_HKEY_CLASSES_ROOT) == 0)
		{
			assertTrue(aux->getName().compare( CONST_HKEY_CLASSES_ROOT) == 0);
			registryNum++;
		}

		else if(aux->getName().compare( CONST_HKEY_CURRENT_USER) == 0)
		{
			assertTrue(aux->getName().compare( CONST_HKEY_CURRENT_USER) == 0);
			registryNum++;
		}

		else if(aux->getName().compare( CONST_HKEY_USERS) == 0)
		{
			assertTrue(aux->getName().compare( CONST_HKEY_USERS) == 0);
			registryNum++;
		}

		else if(aux->getName().compare( CONST_HKEY_CURRENT_CONFIG) == 0)
		{
			assertTrue(aux->getName().compare( CONST_HKEY_CURRENT_CONFIG) == 0);
			registryNum++;
		}
	}

    wmiConnection.close();

	assertTrue(registryNum >= 5);

	cout << "---------- End Standart Registries Test ----------" << endl << endl;
	system("PAUSE");
}

testCase (getLogicalDrives, WMIReaderSuite)
{
	cout << "---------- Start Logical drives Test ----------" << endl;

	while(true)
	{
		WMI_Connection wmiConnection;
		WMI_HResult* whres = initConnection(wmiConnection, false);

		if(FAILED(whres->getHRESULT() ))
		{
			assertTrue(!FAILED(whres->getHRESULT()));
			break;
		}

        WMI_FileDirManager dirMan(&wmiConnection);
		WMI_Result<WMI_Disk> listDrives;
		
        dirMan.getLogicalDrives(listDrives);

        if(!driveActive)
        {
            driveList.clear();
            driveList.push_back("C:");
        }

		int drivesCount = 0;

        for(int i = 0; i < driveList.size(); i++)
		{
			for(int j = 0; j < listDrives.getObjects().size(); j++)
			{
                wstring drive(driveList[i].begin(), driveList[i].end());
                if(drive.compare((listDrives.getObjects()[j])->getName()) == 0)
				{
                    assertTrue(drive.compare((listDrives.getObjects()[j])->getName()) == 0);
					drivesCount++;
				}
			}	
		}

        listDrives.showObjects();
        wmiConnection.close();

        assertTrue(listDrives.getObjects().size() >= 1);
        assertTrue(driveList.size() == drivesCount);

		cout << "---------- End Logical drives Test ----------" << endl << endl;

		system("PAUSE");
		break;
	}
}

testCase(FileDirectoryTest, WMIReaderSuite)
{

	cout << "---------- Start FileDirectory Test ----------" << endl;

	WMI_Connection wmiConnection;
	WMI_HResult* whres = initConnection(wmiConnection, false);

	if(FAILED(whres->getHRESULT() ))
	{
		assertTrue(!FAILED(whres->getHRESULT()));
	}
    delete whres;

	WMI_FileDirManager dirMan(&wmiConnection);
	WMI_Result<WMI_FileDirectory> dirResult;

	if(!fileActive)
	{	
		fileList.clear();
        fileList.push_back(string(FNOTEPAD.begin(), FNOTEPAD.end()));
		fileList.push_back(string(FREGEDIT.begin(), FREGEDIT.end()));
        fileList.push_back(string(FEXPLORER.begin(), FEXPLORER.end()));
        fileList.push_back(string(WINDOWS_DIR.begin(), WINDOWS_DIR.end()));
	}

	wstring fpath;
	wstring fileName;
    int filesCount = 0;

	for (int i = 0; i < fileList.size(); i++)
	{
        wstring path(fileList[i].begin(), fileList[i].end());

		dirMan.getFileDirName(path, fpath, fileName);

		dirMan.getFileDirs(fpath, dirResult);	

		dirResult.showObjects();		

		for(int j = 0; j < dirResult.getObjects()->getContent().size(); j++)
		{
            wstring outFileName;
            dirMan.getFileDirName((dirResult.getObjects()->getContent()[j])->getName(), outFileName);
            transform(outFileName.begin(), outFileName.end(), outFileName.begin(), towlower);
            transform(fileName.begin(), fileName.end(), fileName.begin(), towlower);
            
			if(outFileName.compare(fileName) == 0)
			{
				assertTrue(outFileName.compare(fileName) == 0);
                filesCount++;
			} 
		}
    }

	wmiConnection.close();
    delete whres;

    assertTrue(filesCount == fileList.size());

    cout << "---------- End FileDirectory Test ----------" << endl << endl;
    system("PAUSE");   
}